#include "windowcapture.h"
#include <QDebug>

//using namespace winrt;
using namespace winrt::Windows::Graphics;
using namespace winrt::Windows::Graphics::Capture;
using namespace winrt::Windows::Graphics::DirectX;
using namespace winrt::Windows::Graphics::DirectX::Direct3D11;
using Microsoft::WRL::ComPtr;

// 主线程WinRT初始化标志
static bool g_mainThreadWinRTInitialized = false;
static bool g_mainThreadWinRTInitializationAttempted = false;

struct __declspec(uuid("A9B3D012-3DF2-4EE3-B8D1-8695F457D3C1")) IDirect3DDxgiInterfaceAccess : IUnknown
{
    virtual HRESULT __stdcall GetInterface(GUID const& id, void** object) = 0;
};

// 初始化主线程WinRT（STA）
bool WindowCapture::initializeWinRT()
{
    if (g_mainThreadWinRTInitialized) {
        return true;
    }

    if (g_mainThreadWinRTInitializationAttempted) {
        return false; // 已经尝试过但失败了
    }

    g_mainThreadWinRTInitializationAttempted = true;

    try {
        // 检查WinRT支持
        if (!winrt::Windows::Foundation::Metadata::ApiInformation::IsApiContractPresent(
                L"Windows.Foundation.UniversalApiContract", 8)) {
            qDebug() << "WinRT初始化失败：系统不支持Graphics Capture API";
            return false;
        }

        // 尝试初始化WinRT（主线程使用STA）
        try {
            winrt::init_apartment(winrt::apartment_type::single_threaded);
            qDebug() << "WinRT初始化成功（主线程STA）";
            g_mainThreadWinRTInitialized = true;
            return true;
        } catch (const winrt::hresult_error& err) {
            // 如果已经初始化过，会抛出RPC_E_CHANGED_MODE错误
            if (err.code() == RPC_E_CHANGED_MODE) {
                qDebug() << "WinRT已经初始化过，继续使用";
                g_mainThreadWinRTInitialized = true;
                return true;
            } else {
                qDebug() << "WinRT初始化异常：" << err.code().value << ":" << winrt::to_string(err.message()).c_str();
                return false;
            }
        } catch (...) {
            qDebug() << "WinRT初始化未知异常";
            return false;
        }
    } catch (const winrt::hresult_error& err) {
        qDebug() << "WinRT初始化异常：" << err.code().value << ":" << winrt::to_string(err.message()).c_str();
        return false;
    } catch (...) {
        qDebug() << "WinRT初始化未知异常";
        return false;
    }
}

// 初始化子线程WinRT（MTA）
bool WindowCapture::initializeWinRTForThread()
{
    try {
        // 检查WinRT支持
        if (!winrt::Windows::Foundation::Metadata::ApiInformation::IsApiContractPresent(
                L"Windows.Foundation.UniversalApiContract", 8)) {
            qDebug() << "子线程WinRT初始化失败：系统不支持Graphics Capture API";
            return false;
        }

        // 尝试初始化WinRT（子线程使用MTA）
        try {
            winrt::init_apartment(winrt::apartment_type::multi_threaded);
            qDebug() << "WinRT初始化成功（子线程MTA）";
            return true;
        } catch (const winrt::hresult_error& err) {
            // 如果已经初始化过，会抛出RPC_E_CHANGED_MODE错误
            if (err.code() == RPC_E_CHANGED_MODE) {
                qDebug() << "子线程WinRT已经初始化过，继续使用";
                return true;
            } else {
                qDebug() << "子线程WinRT初始化异常：" << err.code().value << ":" << winrt::to_string(err.message()).c_str();
                return false;
            }
        } catch (...) {
            qDebug() << "子线程WinRT初始化未知异常";
            return false;
        }
    } catch (const winrt::hresult_error& err) {
        qDebug() << "子线程WinRT初始化异常：" << err.code().value << ":" << winrt::to_string(err.message()).c_str();
        return false;
    } catch (...) {
        qDebug() << "子线程WinRT初始化未知异常";
        return false;
    }
}

// 在主线程中创建GraphicsCaptureItem
winrt::Windows::Graphics::Capture::GraphicsCaptureItem WindowCapture::createCaptureItemForWindow(HWND hwnd)
{
    try {
        // 确保主线程WinRT已初始化
        if (!g_mainThreadWinRTInitialized) {
            qDebug() << "主线程WinRT未初始化，无法创建CaptureItem";
            return GraphicsCaptureItem{ nullptr };
        }

        // 检查窗口是否有效
        if (!hwnd || !IsWindow(hwnd)) {
            qDebug() << "无效的窗口句柄";
            return GraphicsCaptureItem{ nullptr };
        }

        if (GetAncestor(hwnd, GA_ROOT) != hwnd) {
            qDebug() << "不是顶级窗口，CreateForWindow 不支持\n";
            return GraphicsCaptureItem{ nullptr };
        }

        if (!IsWindowVisible(hwnd)) {
            qDebug() << "窗口不可见，可能会导致捕获失败\n";
            return GraphicsCaptureItem{ nullptr };
        }

        return createCaptureItem(hwnd);
    } catch (const winrt::hresult_error& err) {
        qDebug() << "创建CaptureItem异常：" << err.code().value << ":" << winrt::to_string(err.message()).c_str();
        return GraphicsCaptureItem{ nullptr };
    } catch (...) {
        qDebug() << "创建CaptureItem未知异常";
        return GraphicsCaptureItem{ nullptr };
    }
}

// 在子线程中捕获（使用已创建的GraphicsCaptureItem）
QImage WindowCapture::captureWithItem(winrt::Windows::Graphics::Capture::GraphicsCaptureItem item)
{
    try {
        // 确保子线程WinRT已初始化
        if (!initializeWinRTForThread()) {
            qDebug() << "子线程WinRT初始化失败";
            return QImage();
        }

        if (!item) {
            qDebug() << "CaptureItem为空";
            return QImage();
        }

        auto d3dDevice = createD3DDevice();
        if (!d3dDevice) {
            qDebug() << "WinRT捕获失败：无法创建D3D设备";
            return QImage();
        }

        auto framePool = Direct3D11CaptureFramePool::Create(
            d3dDevice,
            DirectXPixelFormat::B8G8R8A8UIntNormalized,
            2,  // 使用2个帧缓冲
            item.Size());

        auto session = framePool.CreateCaptureSession(item);

        // 设置无边框捕获（如果支持）
        if (winrt::Windows::Foundation::Metadata::ApiInformation::IsPropertyPresent(
                L"Windows.Graphics.Capture.GraphicsCaptureSession", L"IsBorderRequired")) {
            try {
                GraphicsCaptureAccess::RequestAccessAsync(
                    GraphicsCaptureAccessKind::Borderless).get();
                session.IsBorderRequired(false);
            } catch (...) {
                // 忽略权限错误
            }
        }

        QImage result;
        bool gotFrame = false;
        int frameCount = 0;

        // 设置帧到达回调
        auto frameArrived = framePool.FrameArrived([&](auto&& sender, auto&&) {
            frameCount++;
            if (frameCount > 10) return; // 只处理前10帧

            auto frame = sender.TryGetNextFrame();
            if (frame) {
                auto tex = getTextureFromSurface(frame.Surface());
                if (!tex) return;

                D3D11_TEXTURE2D_DESC desc{};
                tex->GetDesc(&desc);

                // 创建CPU可读的纹理
                D3D11_TEXTURE2D_DESC cpuDesc = desc;
                cpuDesc.Usage = D3D11_USAGE_STAGING;
                cpuDesc.CPUAccessFlags = D3D11_CPU_ACCESS_READ;
                cpuDesc.BindFlags = 0;

                ComPtr<ID3D11Device> device;
                tex->GetDevice(&device);

                ComPtr<ID3D11Texture2D> cpuTex;
                HRESULT hr = device->CreateTexture2D(&cpuDesc, nullptr, &cpuTex);
                if (SUCCEEDED(hr)) {
                    ComPtr<ID3D11DeviceContext> ctx;
                    device->GetImmediateContext(&ctx);
                    ctx->CopyResource(cpuTex.Get(), tex.Get());

                    D3D11_MAPPED_SUBRESOURCE mapped;
                    hr = ctx->Map(cpuTex.Get(), 0, D3D11_MAP_READ, 0, &mapped);
                    if (SUCCEEDED(hr)) {
                        result = QImage(desc.Width, desc.Height, QImage::Format_RGB32);
                        for (UINT y = 0; y < desc.Height; y++) {
                            uchar* dest = result.scanLine(y);
                            uchar* src = reinterpret_cast<uchar*>(mapped.pData) + y * mapped.RowPitch;
                            memcpy(dest, src, desc.Width * 4);
                        }
                        ctx->Unmap(cpuTex.Get(), 0);
                        gotFrame = true;
                    }
                }
            }
        });

        // 启动捕获
        session.StartCapture();

        // 等待帧数据
        for (int i = 0; i < 50 && !gotFrame; i++) {
            Sleep(20);
        }

        // 停止捕获
        try {
            session.Close();
            framePool.Close();
        } catch (...) {
            // 忽略关闭错误
        }

        return result;

    } catch (const winrt::hresult_error& err) {
        qDebug() << "WinRT捕获异常：" << err.code().value << ":" << winrt::to_string(err.message()).c_str();
        return QImage();
    } catch (const std::exception& e) {
        qDebug() << "WinRT捕获异常：" << e.what();
        return QImage();
    } catch (...) {
        qDebug() << "WinRT捕获未知异常";
        return QImage();
    }
}

// 捕获入口（兼容接口）
QImage WindowCapture::captureWindow(HWND hwnd, bool useWinRT)
{
    if (useWinRT) {
        // 在主线程中创建CaptureItem
        auto item = createCaptureItemForWindow(hwnd);
        if (!item) {
            qDebug() << "WinRT捕获失败：无法创建CaptureItem，回退到WinAPI";
            return captureWindowWinAPI(hwnd);
        }
        
        // 在子线程中捕获
        return captureWithItem(item);
    } else {
        return captureWindowWinAPI(hwnd);
    }
}

// WinAPI 捕获
QImage WindowCapture::captureWindowWinAPI(HWND hwnd)
{
    HDC hdcWindow = GetDC(hwnd);
    HDC hdcMemDC = CreateCompatibleDC(hdcWindow);

    RECT rcClient;
    GetClientRect(hwnd, &rcClient);

    HBITMAP hbmScreen = CreateCompatibleBitmap(
        hdcWindow,
        rcClient.right - rcClient.left,
        rcClient.bottom - rcClient.top);

    SelectObject(hdcMemDC, hbmScreen);
    BitBlt(hdcMemDC, 0, 0,
           rcClient.right - rcClient.left,
           rcClient.bottom - rcClient.top,
           hdcWindow, 0, 0, SRCCOPY);

    BITMAP bmpScreen;
    GetObject(hbmScreen, sizeof(BITMAP), &bmpScreen);

    QImage img(bmpScreen.bmWidth, bmpScreen.bmHeight, QImage::Format_RGB32);

    BITMAPINFOHEADER bi = {0};
    bi.biSize = sizeof(BITMAPINFOHEADER);
    bi.biWidth = bmpScreen.bmWidth;
    bi.biHeight = -bmpScreen.bmHeight; // top-down
    bi.biPlanes = 1;
    bi.biBitCount = 32;
    bi.biCompression = BI_RGB;

    GetDIBits(hdcMemDC, hbmScreen, 0,
              (UINT)bmpScreen.bmHeight,
              img.bits(),
              (BITMAPINFO*)&bi, DIB_RGB_COLORS);

    DeleteObject(hbmScreen);
    DeleteDC(hdcMemDC);
    ReleaseDC(hwnd, hdcWindow);

    return img;
}

// 创建 CaptureItem
GraphicsCaptureItem WindowCapture::createCaptureItem(HWND hwnd)
{
    try {
        auto activation_factory = winrt::get_activation_factory<GraphicsCaptureItem>();
        //auto interop_factory = activation_factory.as<IGraphicsCaptureItemInterop>();
        ComPtr<IGraphicsCaptureItemInterop> interop_factory;
        winrt::check_hresult(
            winrt::get_activation_factory<GraphicsCaptureItem>()
            .as<IUnknown>()
            ->QueryInterface(IID_PPV_ARGS(&interop_factory))
        );
        
        GraphicsCaptureItem item = { nullptr };
        //const HRESULT hr = interop_factory->CreateForWindow(
        //    hwnd, 
        //    winrt::guid_of<ABI::Windows::Graphics::Capture::IGraphicsCaptureItem>(),
        //    reinterpret_cast<void**>(winrt::put_abi(item))
        //);
        const HRESULT hr = interop_factory->CreateForWindow(
            hwnd,
            __uuidof(ABI::Windows::Graphics::Capture::IGraphicsCaptureItem),
            reinterpret_cast<void**>(winrt::put_abi(item))
        );
        
        if (FAILED(hr)) {
            qDebug() << "创建CaptureItem失败，HRESULT:" << hr;
            return GraphicsCaptureItem{ nullptr };
        }
        
    return item;
    } catch (const winrt::hresult_error& err) {
        qDebug() << "创建CaptureItem异常：" << err.code().value << ":" << winrt::to_string(err.message()).c_str();
        return GraphicsCaptureItem{ nullptr };
    } catch (...) {
        qDebug() << "创建CaptureItem异常：" << winrt::to_hresult().value;
        return GraphicsCaptureItem{ nullptr };
    }
}

// 创建 D3DDevice
IDirect3DDevice WindowCapture::createD3DDevice()
{
    try {
    ComPtr<ID3D11Device> d3dDevice;
        D3D_FEATURE_LEVEL levels[] = { D3D_FEATURE_LEVEL_11_1, D3D_FEATURE_LEVEL_11_0 };
        D3D_FEATURE_LEVEL featureLevel;
        
        HRESULT hr = D3D11CreateDevice(
            nullptr,
            D3D_DRIVER_TYPE_HARDWARE,
            nullptr,
                      D3D11_CREATE_DEVICE_BGRA_SUPPORT,
            levels,
            ARRAYSIZE(levels),
            D3D11_SDK_VERSION,
            &d3dDevice,
            &featureLevel,
            nullptr
        );
        
        if (FAILED(hr)) {
            qDebug() << "创建D3D11设备失败，HRESULT:" << hr;
            return IDirect3DDevice{ nullptr };
        }

    ComPtr<IDXGIDevice> dxgiDevice;
        hr = d3dDevice->QueryInterface(IID_PPV_ARGS(&dxgiDevice));
        if (FAILED(hr)) {
            qDebug() << "获取DXGI设备失败，HRESULT:" << hr;
            return IDirect3DDevice{ nullptr };
        }

        winrt::com_ptr<::IInspectable> inspectable;
        hr = CreateDirect3D11DeviceFromDXGIDevice(dxgiDevice.Get(), inspectable.put());
        if (FAILED(hr)) {
            qDebug() << "创建WinRT D3D设备失败，HRESULT:" << hr;
            return IDirect3DDevice{ nullptr };
        }

        return inspectable.as<IDirect3DDevice>();
    } catch (const winrt::hresult_error& err) {
        qDebug() << "创建D3D设备异常：" << err.code().value << ":" << winrt::to_string(err.message()).c_str();
        return IDirect3DDevice{ nullptr };
    } catch (...) {
        qDebug() << "创建D3D设备未知异常";
        return IDirect3DDevice{ nullptr };
    }
}

// 从 Surface 提取纹理
ComPtr<ID3D11Texture2D> WindowCapture::getTextureFromSurface(IDirect3DSurface const& surface)
{
    try {
        //auto access = surface.as<Windows::Graphics::DirectX::Direct3D11::IDirect3DDxgiInterfaceAccess>();
        ComPtr<IDirect3DDxgiInterfaceAccess> access;
        winrt::check_hresult(surface.as<IUnknown>()->QueryInterface(IID_PPV_ARGS(&access)));
        ComPtr<ID3D11Texture2D> tex;
        winrt::check_hresult(access->GetInterface(IID_PPV_ARGS(&tex)));
        return tex;
    } catch (const winrt::hresult_error& err) {
        qDebug() << "从Surface提取纹理异常：" << err.code().value << ":" << winrt::to_string(err.message()).c_str();
        return ComPtr<ID3D11Texture2D>{ nullptr };
    } catch (...) {
        qDebug() << "从Surface提取纹理未知异常";
        return ComPtr<ID3D11Texture2D>{ nullptr };
    }
}


// 枚举所有窗口
QMap<HWND, QString> WindowCapture::getAllWindows()
{
    QMap<HWND, QString> windows;

    EnumWindows([](HWND hwnd, LPARAM lParam) -> BOOL {
        if (IsWindowVisible(hwnd)) {
            int length = GetWindowTextLengthW(hwnd);
            if (length > 0) {
                std::wstring title(length + 1, L'\0');
                GetWindowTextW(hwnd, &title[0], length + 1);
                if (!title.empty()) {
                    auto* map = reinterpret_cast<QMap<HWND, QString>*>(lParam);
                    (*map)[hwnd] = QString::fromStdWString(title);
                }
            }
        }
        return TRUE;
    }, reinterpret_cast<LPARAM>(&windows));

    return windows;
}